from fastapi import APIRouter, UploadFile, File, Query, HTTPException, Depends
from fastapi.responses import JSONResponse
import pandas as pd
import io
import os
import uuid
import logging
import tempfile
from typing import List, Dict, Any, Optional

# 配置日志
logger = logging.getLogger(__name__)

router = APIRouter()

# 内存中的Excel缓存
excel_cache: Dict[str, Dict] = {}

# 临时文件目录
TEMP_DIR = tempfile.gettempdir()

class ExcelPreviewResponse:
    """Excel预览响应模型"""
    def __init__(self, headers: List[str], row_count: int, sample_rows: List[Dict[str, Any]], file_id: str):
        self.headers = headers
        self.row_count = row_count
        self.sample_rows = sample_rows
        self.file_id = file_id

class ExcelAskResponse:
    """Excel提问响应模型"""
    def __init__(self, answer: str, context: Dict[str, Any], query: str):
        self.answer = answer
        self.context = context
        self.query = query

def validate_excel_file(file: UploadFile) -> bool:
    """验证Excel文件格式"""
    allowed_extensions = ['.xlsx', '.xls']
    file_extension = os.path.splitext(file.filename)[1].lower()
    
    if file_extension not in allowed_extensions:
        raise HTTPException(status_code=400, detail=f"不支持的文件格式。请上传 .xlsx 或 .xls 文件，当前文件格式: {file_extension}")
    
    # 检查文件大小（限制为10MB）
    file.file.seek(0, 2)
    file_size = file.file.tell()
    file.file.seek(0)  # 重置文件指针
    
    max_size = 10 * 1024 * 1024  # 10MB
    if file_size > max_size:
        raise HTTPException(status_code=400, detail=f"文件大小超过限制。最大允许10MB，当前文件大小: {file_size/1024/1024:.2f}MB")
    
    return True

def read_excel_file(file: UploadFile) -> pd.DataFrame:
    """读取Excel文件为DataFrame"""
    try:
        # 读取文件内容
        contents = file.file.read()
        file.file.seek(0)  # 重置文件指针
        
        # 使用pandas读取Excel
        df = pd.read_excel(io.BytesIO(contents))
        
        # 替换NaN值为空字符串
        df = df.fillna('')
        
        return df
    except Exception as e:
        logger.error(f"读取Excel文件失败: {str(e)}")
        raise HTTPException(status_code=400, detail=f"无法读取Excel文件: {str(e)}")

def cleanup_cache():
    """清理过期的缓存数据，保留最近的20个文件"""
    if len(excel_cache) > 20:
        # 按最后访问时间排序，删除最旧的
        sorted_files = sorted(excel_cache.items(), key=lambda x: x[1]['last_accessed'])
        files_to_remove = len(sorted_files) - 20
        for file_id, _ in sorted_files[:files_to_remove]:
            if file_id in excel_cache:
                del excel_cache[file_id]
                logger.info(f"清理过期的Excel缓存: {file_id}")

@router.post("/preview", response_model=Dict[str, Any])
async def preview_excel(
    file: UploadFile = File(..., description="Excel文件（.xlsx或.xls）")
):
    """
    预览Excel文件内容
    
    - **file**: 要上传的Excel文件
    - 返回表头、总行数、前5行示例数据和文件ID（用于后续提问）
    """
    try:
        # 验证文件
        validate_excel_file(file)
        
        # 读取Excel文件
        df = read_excel_file(file)
        
        # 获取表头
        headers = df.columns.tolist()
        
        # 获取总行数
        row_count = len(df)
        
        # 获取示例数据（最多5行）
        sample_rows = df.head(5).to_dict('records')
        
        # 生成唯一文件ID
        file_id = str(uuid.uuid4())
        
        # 缓存文件数据
        excel_cache[file_id] = {
            'dataframe': df,
            'filename': file.filename,
            'upload_time': pd.Timestamp.now(),
            'last_accessed': pd.Timestamp.now()
        }
        
        # 清理过期缓存
        cleanup_cache()
        
        logger.info(f"预览Excel文件成功: {file.filename}, 总行数: {row_count}")
        
        # 返回结构化响应
        return {
            "headers": headers,
            "row_count": row_count,
            "sample_rows": sample_rows,
            "file_id": file_id,
            "filename": file.filename
        }
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"预览Excel文件失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"预览文件时发生错误: {str(e)}")

@router.post("/ask", response_model=Dict[str, Any])
async def ask_excel(
    query: str = Query(..., description="针对Excel数据的查询问题"),
    file_id: str = Query(..., description="通过预览获取的文件ID"),
    session_id: Optional[str] = Query(None, description="会话ID，用于关联上下文")
):
    """
    基于Excel文件数据回答问题
    
    - **query**: 用户的查询问题
    - **file_id**: 通过preview接口获取的文件ID
    - **session_id**: 可选的会话ID，用于关联查询上下文
    - 返回基于Excel数据的回答
    """
    try:
        # 检查文件ID是否存在
        if file_id not in excel_cache:
            raise HTTPException(status_code=404, detail="Excel文件不存在或已过期，请重新上传预览")
        
        # 获取缓存的DataFrame
        cache_entry = excel_cache[file_id]
        df = cache_entry['dataframe']
        
        # 更新最后访问时间
        cache_entry['last_accessed'] = pd.Timestamp.now()
        
        # 记录查询日志
        logger.info(f"Excel查询: {query} (文件: {cache_entry['filename']}, 会话ID: {session_id})")
        
        # 这里应该集成LLM模型来处理复杂查询
        # 目前使用简单的数据分析示例
        answer = ""
        context = {}
        
        # 简单的分析示例
        try:
            # 示例1: 计算统计信息
            if any(keyword in query.lower() for keyword in ['统计', '汇总', '总数', '平均值']):
                numeric_cols = df.select_dtypes(include=['number']).columns
                if len(numeric_cols) > 0:
                    stats = df[numeric_cols].describe().to_dict()
                    answer = f"已分析Excel数据中的数值列。\n"
                    for col in numeric_cols[:3]:  # 最多显示3列
                        if col in stats:
                            answer += f"- 列 '{col}': 平均值={stats[col]['mean']:.2f}, 最大值={stats[col]['max']:.2f}, 最小值={stats[col]['min']:.2f}\n"
                    context['statistics'] = stats
            
            # 示例2: 查找特定信息
            elif any(keyword in query.lower() for keyword in ['查找', '搜索', '包含']):
                # 这里可以添加更复杂的搜索逻辑
                sample_data = df.head(10).to_dict('records')
                answer = f"这是数据的前10行样本，您可以查看是否包含所需信息。"
                context['sample_data'] = sample_data
            
            # 默认回答
            else:
                answer = "我已成功加载您的Excel文件。请提出更具体的问题，如'计算某列的平均值'或'查找包含特定文本的行'。"
                context['data_overview'] = {
                    'columns': df.columns.tolist(),
                    'row_count': len(df),
                    'numeric_columns': df.select_dtypes(include=['number']).columns.tolist()
                }
        
        except Exception as analysis_error:
            logger.error(f"Excel数据分析失败: {str(analysis_error)}")
            answer = "在分析您的Excel数据时遇到了问题。请尝试提出更具体的问题或检查数据格式是否正确。"
        
        # 返回结构化响应
        return {
            "answer": answer,
            "context": context,
            "query": query,
            "file_id": file_id,
            "filename": cache_entry['filename']
        }
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Excel查询失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"处理查询时发生错误: {str(e)}")

@router.delete("/cache/{file_id}")
async def clear_excel_cache(file_id: str):
    """
    手动清理特定Excel文件的缓存
    
    - **file_id**: 要清理的文件ID
    """
    try:
        if file_id in excel_cache:
            del excel_cache[file_id]
            logger.info(f"已清理Excel缓存: {file_id}")
            return {"message": "缓存已清理"}
        else:
            raise HTTPException(status_code=404, detail="文件缓存不存在")
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"清理Excel缓存失败: {str(e)}")
        raise HTTPException(status_code=500, detail="清理缓存时发生错误")